"""
JWT处理工具模块
实现JWT令牌的创建、验证和刷新功能
"""
from datetime import datetime, timedelta, timezone
from typing import Dict, Any, Optional, Union
from jose import JWTError, jwt
from pydantic import BaseModel

from pkg.settings.config import settings
from pkg.common.constants.auth import TokenType


class TokenPayload(BaseModel):
    """令牌数据模型"""
    sub: Optional[str] = None
    exp: Optional[datetime] = None
    iat: Optional[datetime] = None
    jti: Optional[str] = None
    type: Optional[str] = None


def create_access_token(
    subject: Union[str, Any], 
    expires_delta: Optional[timedelta] = None,
    claims: Dict[str, Any] = None
) -> str:
    """
    创建访问令牌
    :param subject: 令牌主题（通常是用户ID）
    :param expires_delta: 过期时间间隔，如果不提供则使用配置默认值
    :param claims: 额外的声明信息
    :return: JWT令牌字符串
    """
    if expires_delta:
        expire = datetime.now(timezone.utc) + expires_delta
    else:
        expire = datetime.now(timezone.utc) + timedelta(
            minutes=settings.JWT_ACCESS_TOKEN_EXPIRE_MINUTES
        )
    
    # 创建令牌主体
    to_encode = {
        "sub": str(subject),  # 用户ID存储在sub字段中
        "user_id": str(subject),  # 明确添加用户ID字段
        "exp": expire,
        "iat": datetime.now(timezone.utc),
        "type": TokenType.ACCESS
    }
    
    # 添加额外声明
    if claims:
        to_encode.update(claims)
    
    # 加密生成令牌
    encoded_jwt = jwt.encode(
        to_encode, 
        settings.JWT_SECRET_KEY, 
        algorithm=settings.JWT_ALGORITHM
    )
    
    return encoded_jwt


def create_refresh_token(
    subject: Union[str, Any],
    expires_delta: Optional[timedelta] = None
) -> str:
    """
    创建刷新令牌
    :param subject: 令牌主题（通常是用户ID）
    :param expires_delta: 过期时间间隔，如果不提供则使用配置默认值
    :return: JWT令牌字符串
    """
    if expires_delta:
        expire = datetime.now(timezone.utc) + expires_delta
    else:
        expire = datetime.now(timezone.utc) + timedelta(
            days=settings.JWT_REFRESH_TOKEN_EXPIRE_DAYS
        )
    
    # 创建令牌主体
    to_encode = {
        "sub": str(subject),  # 用户ID存储在sub字段中
        "user_id": str(subject),  # 明确添加用户ID字段
        "exp": expire,
        "iat": datetime.now(timezone.utc),
        "jti": f"{datetime.now(timezone.utc).timestamp()}",  # 唯一的令牌ID
        "type": TokenType.REFRESH
    }
    
    # 加密生成令牌
    encoded_jwt = jwt.encode(
        to_encode, 
        settings.JWT_SECRET_KEY, 
        algorithm=settings.JWT_ALGORITHM
    )
    
    return encoded_jwt


def decode_token(token: str) -> TokenPayload:
    """
    解码并验证令牌
    :param token: JWT令牌字符串
    :return: 解码后的令牌数据
    :raises: JWTError 如果令牌无效
    """
    try:
        payload = jwt.decode(
            token, 
            settings.JWT_SECRET_KEY, 
            algorithms=[settings.JWT_ALGORITHM]
        )
        token_data = TokenPayload(**payload)
        
        # 检查令牌是否已过期
        if token_data.exp < datetime.now(timezone.utc):
            raise JWTError("令牌已过期")
            
        return token_data
        
    except JWTError as e:
        raise JWTError(f"无效的令牌: {str(e)}")


def refresh_access_token(refresh_token: str) -> Dict[str, str]:
    """
    使用刷新令牌生成新的访问令牌
    :param refresh_token: 刷新令牌
    :return: 包含新访问令牌的字典
    :raises: JWTError 如果刷新令牌无效或不是刷新类型
    """
    try:
        # 解码刷新令牌
        payload = decode_token(refresh_token)
        
        # 确保是刷新令牌类型
        if payload.type != TokenType.REFRESH:
            raise JWTError("提供的不是刷新令牌")
        
        # 创建新的访问令牌
        new_access_token = create_access_token(subject=payload.sub)
        
        return {
            "access_token": new_access_token,
            "token_type": "bearer"
        }
        
    except JWTError as e:
        raise JWTError(f"刷新令牌失败: {str(e)}")


def verify_token(token: str, token_type: str = TokenType.ACCESS) -> str:
    """
    验证令牌并返回用户ID
    :param token: JWT令牌
    :param token_type: 期望的令牌类型（"access"或"refresh"）
    :return: 用户ID
    :raises: JWTError 如果令牌无效
    """
    payload = decode_token(token)
    
    # 验证令牌类型
    if payload.type != token_type:
        raise JWTError(f"令牌类型错误，期望{token_type}，实际是{payload.type}")
    
    # 返回令牌主题（用户ID）
    return payload.sub
